{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "25bce5eb-8599-40fe-947e-4932cfae8184",
   "metadata": {},
   "source": [
    "# TileDB\n",
    "\n",
    "> [TileDB](https://github.com/TileDB-Inc/TileDB) is a powerful engine for indexing and querying dense and sparse multi-dimensional arrays.\n",
    "\n",
    "> TileDB offers ANN search capabilities using the [TileDB-Vector-Search](https://github.com/TileDB-Inc/TileDB-Vector-Search) module. It provides serverless execution of ANN queries and storage of vector indexes both on local disk and cloud object stores (i.e. AWS S3).\n",
    "\n",
    "More details in:\n",
    "-  [Why TileDB as a Vector Database](https://tiledb.com/blog/why-tiledb-as-a-vector-database)\n",
    "-  [TileDB 101: Vector Search](https://tiledb.com/blog/tiledb-101-vector-search)\n",
    "\n",
    "This notebook shows how to use the `TileDB` vector database."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f45f46f2-7229-4859-9797-30bbead1b8e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install tiledb-vector-search"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f65caa9-8383-409a-bccb-6e91fc8d5e8f",
   "metadata": {},
   "source": [
    "## Basic Example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c96d4fe0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.document_loaders import TextLoader\n",
    "from langchain.embeddings import HuggingFaceEmbeddings\n",
    "from langchain.text_splitter import CharacterTextSplitter\n",
    "from langchain.vectorstores import TileDB\n",
    "\n",
    "raw_documents = TextLoader(\"../../modules/state_of_the_union.txt\").load()\n",
    "text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)\n",
    "documents = text_splitter.split_documents(raw_documents)\n",
    "embeddings = HuggingFaceEmbeddings()\n",
    "db = TileDB.from_documents(\n",
    "    documents, embeddings, index_uri=\"/tmp/tiledb_index\", index_type=\"FLAT\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0a6797c-2bb0-45db-a636-5d2437f7a4c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "query = \"What did the president say about Ketanji Brown Jackson\"\n",
    "docs = db.similarity_search(query)\n",
    "docs[0].page_content"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4c4e06d-6def-44ce-ac9a-4c01673c29a2",
   "metadata": {},
   "source": [
    "### Similarity search by vector"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1eb72610-d451-4158-880c-9f0d45fa5909",
   "metadata": {},
   "outputs": [],
   "source": [
    "embedding_vector = embeddings.embed_query(query)\n",
    "docs = db.similarity_search_by_vector(embedding_vector)\n",
    "docs[0].page_content"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d33588d4-67c2-4bd3-b251-76ae783cbafb",
   "metadata": {},
   "source": [
    "### Similarity search with score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a41e382-0336-4e6d-b2ef-44cc77db2696",
   "metadata": {},
   "outputs": [],
   "source": [
    "docs_and_scores = db.similarity_search_with_score(query)\n",
    "docs_and_scores[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57f930f2-41a0-4795-ad9e-44a33c8f88ec",
   "metadata": {},
   "source": [
    "## Maximal Marginal Relevance Search (MMR)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4790e437-3207-45cb-b121-d857ab5aabd8",
   "metadata": {},
   "source": [
    "In addition to using similarity search in the retriever object, you can also use `mmr` as retriever."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "495754b1-5cdb-4af6-9733-f68700bb7232",
   "metadata": {},
   "outputs": [],
   "source": [
    "retriever = db.as_retriever(search_type=\"mmr\")\n",
    "retriever.get_relevant_documents(query)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e213d957-e439-4bd6-90f2-8909323f5f09",
   "metadata": {},
   "source": [
    "Or use `max_marginal_relevance_search` directly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99d928d0-3b79-4588-925e-32230e12af47",
   "metadata": {},
   "outputs": [],
   "source": [
    "db.max_marginal_relevance_search(query, k=2, fetch_k=10)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
